package com.mezz.bones.framework.gateway.loadbalancer;

import com.mezz.bones.framework.gateway.constant.GatewayConstants;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.*;
import org.springframework.cloud.loadbalancer.core.NoopServiceInstanceListSupplier;
import org.springframework.cloud.loadbalancer.core.ReactorServiceInstanceLoadBalancer;
import org.springframework.cloud.loadbalancer.core.SelectedInstanceCallback;
import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
import org.springframework.http.HttpHeaders;
import reactor.core.publisher.Mono;

import java.util.List;
import java.util.Random;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

/**
 * 基于Nacos元数据版本的自定义负载均衡器
 * 根据请求头中的version参数，在对应版本的服务实例中进行负载均衡
 */
public class MetadataVersionLoadBalancer implements ReactorServiceInstanceLoadBalancer {

    private static final Logger log = LoggerFactory.getLogger(MetadataVersionLoadBalancer.class);

    private final AtomicInteger position = new AtomicInteger(new Random().nextInt(1000));
    private final String serviceId;
    private final ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider;

    public MetadataVersionLoadBalancer(ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider, String serviceId) {
        this.serviceId = serviceId;
        this.serviceInstanceListSupplierProvider = serviceInstanceListSupplierProvider;
    }

    @Override
    public Mono<Response<ServiceInstance>> choose(Request request) {
        ServiceInstanceListSupplier supplier = serviceInstanceListSupplierProvider.getIfAvailable(NoopServiceInstanceListSupplier::new);
        return supplier.get(request).next().map(serviceInstances -> processInstanceResponse(supplier, serviceInstances, request));
    }

    private Response<ServiceInstance> processInstanceResponse(ServiceInstanceListSupplier supplier,
                                                              List<ServiceInstance> serviceInstances,
                                                              Request request) {
        Response<ServiceInstance> serviceInstanceResponse = getInstanceResponse(serviceInstances, request);
        if (supplier instanceof SelectedInstanceCallback && serviceInstanceResponse.hasServer()) {
            ((SelectedInstanceCallback) supplier).selectedServiceInstance(serviceInstanceResponse.getServer());
        }
        return serviceInstanceResponse;
    }

    private Response<ServiceInstance> getInstanceResponse(List<ServiceInstance> instances, Request request) {
        if (instances.isEmpty()) {
            log.warn("没有可用的服务实例: {}", serviceId);
            return new EmptyResponse();
        }

        // 从请求中获取版本信息
        String targetVersion = getVersionFromRequest(request);
        log.info("负载均衡请求 - 服务: {}, 目标版本: {}, 总实例数: {}", serviceId, targetVersion, instances.size());

        // 根据版本过滤实例
        List<ServiceInstance> filteredInstances = filterInstancesByVersion(instances, targetVersion);
        
        if (filteredInstances.isEmpty()) {
            log.warn("没有找到版本为 {} 的服务实例，使用默认负载均衡策略", targetVersion);
            filteredInstances = instances;
        }

        // 在过滤后的实例中进行轮询负载均衡
        ServiceInstance selectedInstance = getInstanceByRoundRobin(filteredInstances);
        
        log.info("选择的服务实例: {}:{}, 版本: {}", 
                selectedInstance.getHost(), 
                selectedInstance.getPort(),
                selectedInstance.getMetadata().get("version"));
        
        return new DefaultResponse(selectedInstance);
    }

    /**
     * 从请求中获取版本信息
     * 默认从请求头中获取
     * 网关 filter可以根据用户特征进行赋值版本信息
     */
    private String getVersionFromRequest(Request request) {
        if (request.getContext() instanceof RequestDataContext) {
            RequestDataContext context = (RequestDataContext) request.getContext();
            HttpHeaders headers = context.getClientRequest().getHeaders();
            
            // 从请求头获取version参数
            String version = headers.getFirst(GatewayConstants.X_ROUTE_VERSION_HEADER_KEY);
            if (version != null && !version.trim().isEmpty()) {
                return version.trim();
            }

        }
        
        // 默认版本
        return "1";
    }

    /**
     * 根据版本过滤服务实例
     */
    private List<ServiceInstance> filterInstancesByVersion(List<ServiceInstance> instances, String targetVersion) {
        return instances.stream()
                .filter(instance -> {
                    String instanceVersion = instance.getMetadata().get("version");
                    // 如果实例没有版本信息，默认为版本1
                    if (instanceVersion == null || instanceVersion.trim().isEmpty()) {
                        instanceVersion = "1";
                    }
                    boolean matches = targetVersion.equals(instanceVersion);
                    log.debug("实例 {}:{} 版本: {}, 目标版本: {}, 匹配: {}", 
                            instance.getHost(), instance.getPort(), instanceVersion, targetVersion, matches);
                    return matches;
                })
                .collect(Collectors.toList());
    }

    /**
     * 轮询算法选择实例
     */
    private ServiceInstance getInstanceByRoundRobin(List<ServiceInstance> instances) {
        int pos = Math.abs(this.position.incrementAndGet());
        ServiceInstance instance = instances.get(pos % instances.size());
        log.debug("轮询选择实例: 位置={}, 实例={}:{}", pos % instances.size(), instance.getHost(), instance.getPort());
        return instance;
    }

}